import sys
from PySide6.QtWidgets import (QApplication, QWidget, QTableWidget, 
                             QTableWidgetItem, QVBoxLayout, QLabel,QHBoxLayout)
from PySide6.QtCore import Qt, QMimeData, QDataStream,QPoint
from PySide6.QtGui import QDrag, QColor, QPixmap, QPainter

class CustomTableWidget(QTableWidget):
    """重写拖放方法的自定义表格控件"""
    def __init__(self, rows=5, cols=3, parent=None):
        super().__init__(rows, cols, parent)
        self.setHorizontalHeaderLabels([f"列 {i+1}" for i in range(cols)])
        self.setSelectionMode(QTableWidget.ExtendedSelection)
        self.setDragDropMode(QTableWidget.DragDrop)
        self.setDefaultDropAction(Qt.MoveAction)
        

        
        # 状态变量
        self.last_highlighted = None  # 记录上一个高亮的单元格
        self.forbidden_cells = {(2, 2)}  # 禁止拖放的单元格(行,列)
        self.number_columns = {1}  # 只接受数字的列(索引)

        # 初始化表格数据
        self._init_table_data()
    def _init_table_data(self):
        """初始化表格数据并标记特殊单元格"""
        for row in range(self.rowCount()):
            for col in range(self.columnCount()):
                item = QTableWidgetItem(f"数据 ({row+1},{col+1})")
                item.setTextAlignment(Qt.AlignCenter)
                
                # 标记禁止拖放的单元格
                if (row, col) in self.forbidden_cells:
                    item.setBackground(QColor(255, 220, 220))
                    item.setToolTip("禁止拖放")
                # 标记数字列
                elif col in self.number_columns:
                    item.setBackground(QColor(220, 255, 220))
                    item.setToolTip("仅接受数字")
                    
                self.setItem(row, col, item)

    # ------------------------------
    # 源控件方法: 处理拖动数据
    # ------------------------------
    def startDrag(self, supported_actions):
        """重写开始拖动方法，自定义拖动数据和视觉反馈"""
        selected_items = self.selectedItems()
        if not selected_items:
            return super().startDrag(supported_actions)
        
        # 过滤掉禁止拖动的单元格
        valid_items = [
            item for item in selected_items
            if (item.row(), item.column()) not in self.forbidden_cells
        ]
        if not valid_items:
            return
        
        # 创建MIME数据
        mime_data = QMimeData()
        
        # 1. 添加Qt默认表格格式(便于表格间拖放)
        data_stream = QDataStream()
        data_stream.writeInt32(len(valid_items))
        
        for item in valid_items:
            # 写入行列索引
            data_stream.writeInt32(item.row())
            data_stream.writeInt32(item.column())
            
            # 写入角色-值对(显示文本和背景色)
            data_stream.writeInt32(2)  # 两个角色
            # 显示文本角色
            data_stream.writeInt32(Qt.DisplayRole)
            data_stream.writeQVariant(item.text())
            # 背景色角色
            data_stream.writeInt32(Qt.BackgroundRole)
            data_stream.writeQVariant(item.background())
        
        # mime_data.setData(
        #     "application/x-qabstractitemmodeldatalist",
        #     data_stream.device().data()
        # )
        
        # 2. 添加文本格式(兼容其他控件)
        text_data = [f"{item.row()},{item.column()},{item.text()}" for item in valid_items]
        mime_data.setText(";".join(text_data))
        
        # 创建拖动对象
        drag = QDrag(self)
        drag.setMimeData(mime_data)
        
        # 设置拖动图标
        drag.setPixmap(self._create_drag_pixmap(item.text()))
        drag.setHotSpot(QPoint(-10, 0))
        
        # 执行拖动并处理结果
        result = drag.exec_(supported_actions)
        if result == Qt.MoveAction:
            # 移动操作: 删除源单元格
            for item in valid_items:
                self.takeItem(item.row(), item.column())

    def _create_drag_pixmap(self, dataText):
        """创建自定义拖动图标"""
        pixmap = QPixmap(80, 25)
        pixmap.fill(QColor(255, 255, 255,200))
        
        painter = QPainter(pixmap)
        painter.setPen(QColor(0, 0, 0))
        painter.drawText(pixmap.rect(), Qt.AlignCenter, f"{dataText}")
        painter.drawRect(0, 0, pixmap.width()-1, pixmap.height()-1)
        painter.end()
        
        return pixmap

    # ------------------------------
    # 目标控件方法: 处理拖入数据
    # ------------------------------
    def dragEnterEvent(self, event):
        """拖入时验证数据"""
        if event.source() != self and (
            event.mimeData().hasText() or 
            event.mimeData().hasFormat("application/x-qabstractitemmodeldatalist")
        ):
            event.acceptProposedAction()
        else:
            event.ignore()

    def dragMoveEvent(self, event):
        """拖动过程中提供视觉反馈和位置验证"""
        # 清除上一次高亮
        self._clear_highlight()
        
        # 获取目标位置
        pos = event.pos()
        index = self.indexAt(pos)
        
        if not index.isValid():
            event.ignore()
            return
            
        row, col = index.row(), index.column()
        
        # 检查是否为禁止区域
        if (row, col) in self.forbidden_cells:
            event.ignore()
            return
            
        # 高亮目标单元格
        self._highlight_cell(row, col)
        
        event.acceptProposedAction()

    def dragLeaveEvent(self, event):
        """拖动离开时清除高亮"""
        self._clear_highlight()
        super().dragLeaveEvent(event)

    def dropEvent(self, event):
        """处理最终放下的数据"""
        # 清除高亮
        self._clear_highlight()
        
        # 获取目标位置
        pos = event.pos()
        index = self.indexAt(pos)
        target_row, target_col = 0, 0
        
        if index.isValid():
            target_row, target_col = index.row(), index.column()
        else:
            # 拖到空白处，放到最后一行
            target_row = self.rowCount()
            target_col = 0
            self.insertRow(target_row)
        
        # 解析并处理数据
        mime_data = event.mimeData()
        if mime_data.hasFormat("application/x-qabstractitemmodeldatalist"):
            self._process_model_data(mime_data, target_row, target_col)
        elif mime_data.hasText():
            self._process_text_data(mime_data.text(), target_row, target_col)
            
        event.acceptProposedAction()

    # ------------------------------
    # 辅助方法
    # ------------------------------
    def _highlight_cell(self, row, col):
        """高亮目标单元格"""
        item = self.item(row, col) or QTableWidgetItem()
        self.last_highlighted = {
            "row": row,
            "col": col,
            "bg": item.background()
        }
        item.setBackground(QColor(255, 245, 157))
        if not self.item(row, col):
            self.setItem(row, col, item)

    def _clear_highlight(self):
        """清除单元格高亮"""
        if self.last_highlighted:
            row = self.last_highlighted["row"]
            col = self.last_highlighted["col"]
            item = self.item(row, col)
            if item:
                item.setBackground(self.last_highlighted["bg"])
            self.last_highlighted = None

    def _process_text_data(self, text, start_row, start_col):
        """处理文本格式数据"""
        items = text.split(';')
        current_row, current_col = start_row, start_col
        
        for item_data in items:
            parts = item_data.split(',', 2)
            if len(parts) != 3:
                continue
            _, _, text = parts
            
            # 检查数字列
            if current_col in self.number_columns and not text.replace('.', '', 1).isdigit():
                text = f"[{text}]"  # 标记无效数字
            
            self._insert_item(current_row, current_col, text)
            
            # 处理列溢出
            current_col += 1
            if current_col >= self.columnCount():
                current_row += 1
                current_col = 0
                if current_row >= self.rowCount():
                    self.insertRow(current_row)

    def _process_model_data(self, mime_data, start_row, start_col):
        """处理Qt表格格式数据"""
        data = mime_data.data("application/x-qabstractitemmodeldatalist")
        stream = QDataStream(data, QDataStream.ReadOnly)
        
        item_count = stream.readInt32()
        current_row, current_col = start_row, start_col
        
        for _ in range(item_count):
            # 源行列索引(这里不需要使用)
            stream.readInt32()
            stream.readInt32()
            
            role_count = stream.readInt32()
            item_data = {}
            
            for _ in range(role_count):
                role = stream.readInt32()
                value = stream.readQVariant()
                item_data[role] = value
            
            # 提取显示文本和样式
            text = item_data.get(Qt.DisplayRole, "未知数据")
            background = item_data.get(Qt.BackgroundRole, QColor(255, 255, 255))
            
            # 插入单元格
            item = self._insert_item(current_row, current_col, text)
            item.setBackground(background)
            
            # 处理列溢出
            current_col += 1
            if current_col >= self.columnCount():
                current_row += 1
                current_col = 0
                if current_row >= self.rowCount():
                    self.insertRow(current_row)

    def _insert_item(self, row, col, text):
        """插入单元格并返回创建的项目"""
        item = QTableWidgetItem(text)
        item.setTextAlignment(Qt.AlignCenter)
        
        # 数字列特殊样式
        if col in self.number_columns:
            item.setForeground(QColor(139, 0, 0) if '[' in text else QColor(0, 100, 0))
        
        self.setItem(row, col, item)
        return item

class DemoWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("QTableWidget 重写拖放方法示例")
        self.setGeometry(100, 100, 1000, 600)
        
        layout = QVBoxLayout()
        
        info_label = QLabel("""
        拖放功能说明：<br>
        - 红色单元格：禁止拖放<br>
        - 绿色单元格：仅接受数字（非数字会标记）<br>
        - 拖动时会高亮目标位置<br>
        - 支持表格间互拖和表格内拖放
        """)
        layout.addWidget(info_label)
        
        # 创建两个表格用于测试
        table_layout = QHBoxLayout()
        self.table1 = CustomTableWidget()
        self.table2 = CustomTableWidget()
        table_layout.addWidget(self.table1)
        table_layout.addWidget(self.table2)
        
        layout.addLayout(table_layout)
        self.setLayout(layout)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = DemoWindow()
    window.show()
    sys.exit(app.exec())
    